package org.example.openAi;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author HuangBenben
 */
public class AiUtil {

    private static final ClientV4 client = new ClientV4
            .Builder("51ca0c08adf2f23443414b8dd87a6fc9.vGRzvdzXWC7iogco")
            .networkConfig(1000 * 20, 1000 * 20, 1000 * 20, 1000 * 20, TimeUnit.MILLISECONDS)
            .build();

    private static final ObjectMapper mapper = new ObjectMapper();



    /**
     * 同步调用
     */
    public static String invoke(String question) {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), question);
        messages.add(chatMessage);
        String requestId = UUID.randomUUID().toString();

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();

        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        try {
            return mapper.writeValueAsString(invokeModelApiResp);
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        return "";
    }

    public static void sseInvoke(String msg) {
        // 创建一个消息列表，将用户消息添加进去
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), msg);
        messages.add(chatMessage);
        String requestId = UUID.randomUUID().toString();
        // 构建请求对象，配置为流式请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)      // 设置使用的模型
                .stream(Boolean.TRUE)                // 启用流式返回
                .messages(messages)                  // 设置消息
                .requestId(requestId)                // 设置请求ID
                .build();

        // 调用 API 获取流式响应
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);

        // 判断 API 请求是否成功
        if (sseModelApiResp.isSuccess()) {
            // 使用 AtomicBoolean 来控制第一次输出
            AtomicBoolean isFirst = new AtomicBoolean(true);

            // 将流数据映射到累加器（即接收器）
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable(), chatMessage)
                    .doOnNext(accumulator -> {
                        // 处理流数据
                        if (isFirst.getAndSet(false)) {
                            System.out.print("Response: ");  // 打印初始响应
                        }

                        // 处理 tool_calls，如果有
                        if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                            try {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            } catch (JsonProcessingException e) {
                                e.printStackTrace();  // 错误处理
                            }
                        }

                        // 打印返回的具体内容
                        if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                            System.out.print(accumulator.getDelta().getContent());
                        }
                    })
                    .doOnComplete(System.out::println)  // 完成时打印
                    .lastElement()  // 获取最后的元素
                    .blockingGet();  // 阻塞并获取结果

            // 创建一个新的 Choice 对象用于返回结果
            Choice choice = new Choice();
            choice.setFinishReason(chatMessageAccumulator.getChoice().getFinishReason());
            choice.setIndex(0L);
            choice.setDelta(chatMessageAccumulator.getDelta());

            // 将 Choice 放入一个列表中
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);

            // 创建返回的数据对象
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());

            // 清空流并设置返回数据
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
    }

    private static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable, ChatMessage chatMessage) {
        return flowable
                .map(modelData -> {
                    // 提取 ModelData 中的必要属性
                    Delta delta = modelData.getChoices().get(0).getDelta();
                    Choice choice = modelData.getChoices().get(0);
                    Usage usage = modelData.getUsage();
                    String id = modelData.getId();
                    long created = modelData.getCreated();

                    // 使用带参数的构造函数创建 ChatMessageAccumulator 对象
                    ChatMessageAccumulator accumulator = new ChatMessageAccumulator(delta, chatMessage, choice, usage, created, id);
                    return accumulator;
                });
    }


    public static void sseInvokeNew(String msg) {
        // 创建一个消息列表，将用户消息添加进去
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), msg);
        messages.add(chatMessage);
        String requestId = UUID.randomUUID().toString();
        // 构建请求对象，配置为流式请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)      // 设置使用的模型
                .stream(Boolean.TRUE)                // 启用流式返回
                .messages(messages)                  // 设置消息
                .requestId(requestId)                // 设置请求ID
                .build();

        // 调用 API 获取流式响应
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);

        // 判断 API 请求是否成功
        if (sseModelApiResp.isSuccess()) {
            // 使用 AtomicBoolean 来控制第一次输出
            AtomicBoolean isFirst = new AtomicBoolean(true);

            // 将流数据映射到累加器（即接收器）
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable(), chatMessage)
                    .doOnNext(accumulator -> {
                        // 处理流数据
                        if (isFirst.getAndSet(false)) {
                            System.out.print("Response: ");  // 打印初始响应
                        }

                        // 处理 tool_calls，如果有
                        if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                            try {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            } catch (JsonProcessingException e) {
                                e.printStackTrace();  // 错误处理
                            }
                        }

                        // 打印返回的具体内容
                        if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                            System.out.print(accumulator.getDelta().getContent());
                        }
                    })
                    .doOnComplete(System.out::println)  // 完成时打印
                    .lastElement()  // 获取最后的元素
                    .blockingGet();  // 阻塞并获取结果

            // 创建一个新的 Choice 对象用于返回结果
            Choice choice = new Choice();
            choice.setFinishReason(chatMessageAccumulator.getChoice().getFinishReason());
            choice.setIndex(0L);
            choice.setDelta(chatMessageAccumulator.getDelta());

            // 将 Choice 放入一个列表中
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);

            // 创建返回的数据对象
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());

            // 清空流并设置返回数据
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
    }


}
